{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获得图片路径列表，并且划分训练集和测试集\n",
    "all_parent_path = []\n",
    "all_data_list = []\n",
    "all_label_list = []\n",
    "\n",
    "source_path = r\"E:\\jupyter-notebook\\TeacherWork\\data\\shapes\"\n",
    "for p in os.listdir(source_path):\n",
    "    all_parent_path.append(os.path.join(source_path, p))\n",
    "\n",
    "#获取所有的图像和对应标签的列表\n",
    "for img_parent_path in all_parent_path:\n",
    "    for img in os.listdir(img_parent_path):\n",
    "        all_data_list.append(os.path.join(img_parent_path, img))\n",
    "        all_label_list.append(img_parent_path.split(\"\\\\\")[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "E:\\jupyter-notebook\\TeacherWork\\data\\shapes\\triangles\\drawing(1).png\n",
      "triangles\n"
     ]
    }
   ],
   "source": [
    "print(all_data_list[200])\n",
    "print(all_label_list[200])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n",
      "三角形\n"
     ]
    }
   ],
   "source": [
    "# img表示输入的图片，即为需要进行形状判断的图片\n",
    "img = all_data_list[200]\n",
    "frame = cv2.imread(img)\n",
    "h, w, ch = frame.shape\n",
    "result = np.zeros((h, w, ch), dtype=np.uint8)\n",
    "# 二值化图像\n",
    "gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)\n",
    "#cv2.imshow(\"input image\", frame)\n",
    "#print(cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE))\n",
    "contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "for cnt in range(len(contours)):\n",
    "            # 提取与绘制轮廓\n",
    "            cv2.drawContours(result, contours, cnt, (0, 255, 0), 2)\n",
    "\n",
    "            # 轮廓逼近\n",
    "            epsilon = 0.05 * cv2.arcLength(contours[cnt], True)\n",
    "            approx = cv2.approxPolyDP(contours[cnt], epsilon, True)\n",
    "            #print(approx)\n",
    "\n",
    "            # 分析几何形状\n",
    "            corners = len(approx)\n",
    "            print(corners)\n",
    "            shape_type = \"\"\n",
    "            if corners == 3:\n",
    "                shape_type = \"三角形\"\n",
    "                print(shape_type)\n",
    "            if corners == 4:\n",
    "                shape_type = \"矩形\"\n",
    "                print(shape_type)\n",
    "            if corners > 4:\n",
    "                shape_type = \"圆形\"\n",
    "                print(shape_type)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.8 ('TorchForGpu')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "8cdde082121d1919a4a986b8b43c23274ab05bcf17a9fec37e55041b71ca8bca"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
